#pragma once

#include <stdlib.h>
#include <gflags/gflags.h>

#include <doca_log.h>
#include <infiniband/verbs.h>
#include <infiniband/mlx5_api.h>
#include <libflexio/flexio.h>
#include <numa.h>
#include <butil/logging.h>

#include "wrapper_flexio.hpp"
#include "common_cross.h"
#include "common_host.hpp"

#include "common.hpp"
#include "feature_cross.h"
#include "cpu_feature_operator.h"
#include "timer.h"

#include "dpa_common.h"
#include <mutex>

#define DEFAULT_WINDOWS_SIZE 1024 * 1024 * 1024
// default dpu thread_num
#define DEFAULT_THREAD_NUM 16

#define DEFAULT_BATCH_SIZE 1

#if defined(__aarch64__)
DEFINE_string(rdma_device, "mlx5_bond_0", "rdma_device_name");
#else
DEFINE_string(rdma_device, "mlx5_3", "rdma_device_name");
#endif

class DpaUnit{
public:
    int get_thread_num() {return active_thread_num;}
    /* malloc在dma进行，这里只需要注册，映射好 */
    virtual void register_host_memory(void * ptr, size_t size) = 0;

    virtual void submit_task(uint64_t ddata, int execute_thread_num = 1) = 0;

    virtual void init_dpa_thread_info(DpaThreadInfo & T) = 0;

    virtual void poll_thread() = 0;

    virtual size_t get_execute_cycle() = 0;

public:
    int active_thread_num;
};

class DpaFlexioUnit : public DpaUnit {
public:
    DpaFlexioUnit() {
        this->active_thread_num = MAX_DPA_THREAD_NUM;
        this->device_name = FLAGS_rdma_device;
        app = feature_device;
        func = sequence_preprocess;
        init_func = dpa_execute_unit_init;

        this->context_init(this->app);
        this->log_init(0);
        this->command_queue_init();

        LOG(INFO) << "finish DpaFlexioUnit init";
    }

    void context_init(flexio_app *app_device) {
        ctx = new FLEX::Context(device_name);
        ctx->alloc_pd();
        ctx->create_process(app_device);
        ctx->create_window();
        ctx->generate_flexio_uar();
    }

    void log_init(int num) { //给多少个函数进行log输出
        ctx->print_init("mt_memory", num);
    }

    void command_queue_init() {
        cmdq = new FLEX::CommandQ(ctx->get_process(), active_thread_num, DEFAULT_BATCH_SIZE);
    }

    /* malloc在dma进行，这里只需要注册，映射好 */
    void register_host_memory(void * ptr, size_t size) override{
        ctx->reg_mr(ptr, (size + 63) & (~63));
        LOG(INFO) << "successful finish register host memory";
    }

    /* 实验证明分配多个lkey不成功，还是只能用一大片内存，同一个window */
    uint32_t get_lkey() {
        return ctx->get_mr(0)->lkey;
    }

    uint32_t get_window_id() {
        return ctx->get_window_id();
    }

    void submit_task(uint64_t ddata, int execute_thread_num) override{
        (void) execute_thread_num;
        cmdq->add_task(func, ddata);
    }

    void init_addr(void * input_mem_addr, void * output_mem_addr) {
        ptr_mask.input_mem_addr = input_mem_addr;
        ptr_mask.output_mem_addr = output_mem_addr;
        flexio_host2dev_memcpy(ctx->get_process(), &ptr_mask, sizeof(external_ptr), ptr_info);
    }

    void init_dpa_thread_info(DpaThreadInfo & T) override {
        T.window_id       = this->get_window_id();
        T.input_mem_mkey  = this->get_lkey();
        T.output_mem_mkey = this->get_lkey();
        uint64_t rpc_ret_val;
        flexio_uintptr_t dev_config_data = ctx->move_to_dev(T);
        Assert(flexio_process_call(ctx->get_process(), &dpa_execute_unit_init, &rpc_ret_val,
            dev_config_data) == FLEXIO_STATUS_SUCCESS);

        // test heap
        std::call_once(once_flag, [&] {
        // 假设 external_ptr 是已定义类型，否则需正确定义
            this->ptr_info = ctx->move_to_dev(ptr_mask);
            flexio_process_call(ctx->get_process(), &set_addr, &rpc_ret_val, ptr_info);
        }); 
    }

    void poll_thread() {
        // 不知道这个commandQ的pending还是running对过程有何影响？需要pending才能加task么？
        cmdq->run();
        // TODO: 把poll wait中的while循环中的sleep加上，不然一直占用资源。
        auto time_us = cmdq->wait_run(300); // 100 seconds
        (void)time_us;
    }

    size_t get_execute_cycle() override {
        size_t execute_cycle = 0;
        assert(flexio_process_call(ctx->get_process(), get_max_thread_time, &execute_cycle) == FLEXIO_STATUS_SUCCESS);
        return execute_cycle;
    }

    ~DpaFlexioUnit() {
        
    }


public:
    flexio_func_t * func;
    flexio_func_t * init_func;
    struct flexio_app * app;
    std::string device_name;
    FLEX::Context * ctx;
    FLEX::CommandQ * cmdq;

    // test heap
    external_ptr ptr_mask;
    flexio_uintptr_t ptr_info;
    std::once_flag once_flag;
};

class DpaDocasdkUnit : public DpaUnit{
public:
    DpaDocasdkUnit() {
        this->active_thread_num = 128; //TODO: max_thread_num_per_kernel = 128
        	
        doca_error_t result;

        strcpy(cfg.pf_device_name, DEVICE_DEFAULT_NAME);
        strcpy(cfg.rdma_device_name, DEVICE_DEFAULT_NAME);

        result = doca_log_backend_create_standard();
        result = doca_log_backend_create_with_file_sdk(stderr, &sdk_log);
        result = doca_log_backend_set_sdk_level(sdk_log, DOCA_LOG_LEVEL_WARNING);

        result = allocate_dpa_resources(&cfg, &resources);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "Failed to Allocate DPA Resources";
        }

        LOG(INFO) << "finish dpa allocate";

        unsigned int value = 0;
        doca_dpa_get_max_threads_per_kernel(resources.pf_dpa_ctx, &value);
        LOG(INFO) << "max thread per kernel = " << value;

        unsigned long long max_time = 0;
        doca_dpa_get_kernel_max_run_time(resources.pf_dpa_ctx, &max_time);
        LOG(INFO) << "max execute time per kernel = " << max_time;

        comp_event = NULL;
        comp_event_val = 10;

        result = create_doca_dpa_completion_sync_event(resources.pf_dpa_ctx,
                                resources.pf_doca_device,
                                &comp_event,
                                NULL);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "sync_event_completion failed";
        }

        func = docasdkprocess;
        init_func = dpa_execute_unit_init;

        LOG(INFO) << "finish sync event register";
    }

    void register_host_memory(void * ptr, size_t size) override {
        void * send_ptr = ptr;
        host_mmap_obj.mmap_type = MMAP_TYPE_CPU;
        host_mmap_obj.doca_dpa = resources.pf_dpa_ctx;
        host_mmap_obj.doca_device = resources.pf_doca_device;
        host_mmap_obj.permissions = DOCA_ACCESS_FLAG_LOCAL_READ_WRITE | DOCA_ACCESS_FLAG_RDMA_WRITE |
                        DOCA_ACCESS_FLAG_RDMA_READ | DOCA_ACCESS_FLAG_RDMA_ATOMIC | DOCA_ACCESS_FLAG_PCI_READ_WRITE;
        host_mmap_obj.memrange_addr = send_ptr;
        host_mmap_obj.memrange_len = ( (size + 63) & (~63) );

        doca_error_t cur_doca_err = doca_mmap_obj_init(&host_mmap_obj);
        if (cur_doca_err != DOCA_SUCCESS) {printf ("register_host_memory error!\n"); fflush(stdout);}
    }

    void submit_task(uint64_t ddata, int execute_thread_num) override {
        doca_error_t result;
        result = doca_dpa_kernel_launch_update_set(resources.pf_dpa_ctx,
						   NULL,
						   0,
						   comp_event,
						   comp_event_val,
						   execute_thread_num,
						   func,
						   host_mmap_obj.dpa_mmap_handle, 
						   (uint64_t)ddata);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << "kernel_launch failed";
        }
    }

    void init_dpa_thread_info(DpaThreadInfo & T) override {
        uint64_t rpc_ret_val;
        doca_dpa_dev_uintptr_t dev_addr; 
        doca_error_t doca_err = DOCA_SUCCESS;
        doca_err = doca_dpa_mem_alloc(resources.pf_dpa_ctx, sizeof(T), &dev_addr);
        doca_dpa_h2d_memcpy(resources.pf_dpa_ctx, dev_addr, &T, sizeof(T) );
        if (doca_err != DOCA_SUCCESS) {
            LOG(INFO) << "Function doca_dpa_rpc failed";
        }

        doca_err = doca_dpa_rpc(resources.pf_dpa_ctx,
				init_func,
				&rpc_ret_val,
				dev_addr);
        if (doca_err != DOCA_SUCCESS) {
            LOG(INFO) << "Function doca_dpa_rpc failed";
        }
    }

    void poll_thread() override {
        doca_error_t result;
        result = doca_sync_event_wait_gt(comp_event, comp_event_val - 1, SYNC_EVENT_MASK_FFS);
        if (result != DOCA_SUCCESS) {
            LOG(INFO) << ("poll_wait error");
        }
    }

    size_t get_execute_cycle() override {
        return 0;
    }

public:
    struct dpa_config cfg;
    struct dpa_resources resources;
    struct doca_mmap_obj host_mmap_obj;
    struct doca_sync_event * comp_event;
    struct doca_log_backend * sdk_log;
    uint64_t comp_event_val;

    doca_dpa_func_t * func;
    doca_dpa_func_t * init_func;
};

/* test if there some way efficient launch thread */

class DpaTestUnit : public DpaUnit {
public:
    DpaTestUnit() {
        this->active_thread_num = MAX_DPA_THREAD_NUM;
        this->device_name = FLAGS_rdma_device;
        app = feature_device;
        func = test_preprocess;
        init_func = dpa_execute_unit_init;

        this->context_init(this->app);
        this->log_init(0);
        this->command_queue_init();

        LOG(INFO) << "finish DpaFlexioUnit init";
    }

    void context_init(flexio_app *app_device) {
        ctx = new FLEX::Context(device_name);
        ctx->alloc_pd();
        ctx->create_process(app_device);
        ctx->create_window();
        ctx->generate_flexio_uar();
    }

    void log_init(int num) { //给多少个函数进行log输出
        ctx->print_init("mt_memory", num);
    }

    void command_queue_init() {
        cmdq = new FLEX::CommandQ(ctx->get_process(), active_thread_num, DEFAULT_BATCH_SIZE);
    }

    /* malloc在dma进行，这里只需要注册，映射好 */
    void register_host_memory(void * ptr, size_t size) override{
        ctx->reg_mr(ptr, (size + 63) & (~63));
        LOG(INFO) << "successful finish register host memory";
    }

    /* 实验证明分配多个lkey不成功，还是只能用一大片内存，同一个window */
    uint32_t get_lkey() {
        return ctx->get_mr(0)->lkey;
    }

    uint32_t get_window_id() {
        return ctx->get_window_id();
    }

    void submit_task(uint64_t ddata, int execute_thread_num) override{
        (void) execute_thread_num;
        ( (volatile uint64_t*)ctrl_ptr)[1] = 0;
        ( (volatile uint64_t*)ctrl_ptr)[0] = ddata;
    }

    void init_dpa_thread_info(DpaThreadInfo & T) override {
        T.window_id       = this->get_window_id();
        T.input_mem_mkey  = this->get_lkey();
        T.output_mem_mkey = this->get_lkey();
        uint64_t rpc_ret_val;
        flexio_uintptr_t dev_config_data = ctx->move_to_dev(T);
        Assert(flexio_process_call(ctx->get_process(), &dpa_execute_unit_init, &rpc_ret_val,
            dev_config_data) == FLEXIO_STATUS_SUCCESS);
        // FIXME: 可能单独一个函数用来初始化
        this->thread_num = T.thread_num;
    }

    void init_control_info(void * ptr) {
        ctrl_ptr = (uint64_t *) ptr;
        uint64_t rpc_ret_val;

        ThreadControl T;
        T.window_id = this->get_window_id();
        T.mem_mkey  = this->get_lkey();
        T.mem_addr  = ptr;

        flexio_uintptr_t dev_config_data = ctx->move_to_dev(T);
        Assert(flexio_process_call(ctx->get_process(), &threadcontrol_init, &rpc_ret_val,
            dev_config_data) == FLEXIO_STATUS_SUCCESS);

        std::call_once(once_flag, [&] {
            ( (volatile uint64_t*)ctrl_ptr)[0] = 0;
            for (int i = 0; i < this->thread_num; ++i) {
                cmdq->add_task(func, i);
            }
            cmdq->run();
            sleep(3);
        });
    }

    void poll_thread() {
        // while (1) {
        //     printf ("%lu\n", ctrl_ptr[1]);
        //     if (ctrl_ptr[1] == 254) break;
        //     sleep(2);
        // }
        // auto time_us = cmdq->wait_run(300); // 100 seconds
        //  while (!flexio_cmdq_is_empty(cmdq->cmd_q)) {
        while (1) {
            // sleep(2);
            if (( (volatile uint64_t*) ctrl_ptr)[1] == (uint64_t)(this->thread_num) ) {
                // printf ("ctrl[1] = %lu\n", ( (volatile uint64_t*) ctrl_ptr)[1]); 
                // fflush(stdout);
                ( (volatile uint64_t*) ctrl_ptr)[1] = 0;
                break;
            }
            // usleep(step_us);
            // useconds += step_us;
            // Assert(useconds < time_out_seconds * 1e6);
            // break;
        }
        // (void)time_us;
    }

    size_t get_execute_cycle() override {
        size_t execute_cycle = 0;
        assert(flexio_process_call(ctx->get_process(), get_max_thread_time, &execute_cycle) == FLEXIO_STATUS_SUCCESS);
        return execute_cycle;
    }

    ~DpaTestUnit() {
        
    }


public:
    flexio_func_t * func;
    flexio_func_t * init_func;
    struct flexio_app * app;
    std::string device_name;
    FLEX::Context * ctx;
    FLEX::CommandQ * cmdq;
    uint64_t * ctrl_ptr;
    std::once_flag once_flag;
    int thread_num = 0;
};